{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) MONAI Consortium  \n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    "you may not use this file except in compliance with the License.  \n",
    "You may obtain a copy of the License at  \n",
    "&nbsp;&nbsp;&nbsp;&nbsp;http://www.apache.org/licenses/LICENSE-2.0  \n",
    "Unless required by applicable law or agreed to in writing, software  \n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    "See the License for the specific language governing permissions and  \n",
    "limitations under the License.\n",
    "\n",
    "# 3D Multi-organ Segmentation with UNETR  (BTCV Challenge)\n",
    "# PyTorch Lightning Tutorial\n",
    "\n",
    "\n",
    "This tutorial demonstrates how MONAI can be used in conjunction with PyTorch Lightning framework to construct a training workflow of UNETR on multi-organ segmentation task using the BTCV challenge dataset.\n",
    "\n",
    "![image](../figures/btcv_transformer.png)\n",
    "\n",
    "And it contains the following features:\n",
    "1. Transforms for dictionary format data.\n",
    "2. Define a new transform according to MONAI transform API.\n",
    "3. Load Nifti image with metadata, load a list of images and stack them.\n",
    "4. Randomly adjust intensity for data augmentation.\n",
    "5. Cache IO and transforms to accelerate training and validation.\n",
    "6. 3D UNETR model, Dice loss function, Mean Dice metric for multi-oorgan segmentation task.\n",
    "\n",
    "The dataset comes from https://www.synapse.org/#!Synapse:syn3193805/wiki/217752.  \n",
    "\n",
    "Under Institutional Review Board (IRB) supervision, 50 abdomen CT scans of were randomly selected from a combination of an ongoing colorectal cancer chemotherapy trial, and a retrospective ventral hernia study. The 50 scans were captured during portal venous contrast phase with variable volume sizes (512 x 512 x 85 - 512 x 512 x 198) and field of views (approx. 280 x 280 x 280 mm3 - 500 x 500 x 650 mm3). The in-plane resolution varies from 0.54 x 0.54 mm2 to 0.98 x 0.98 mm2, while the slice thickness ranges from 2.5 mm to 5.0 mm. \n",
    "\n",
    "Target: 13 abdominal organs including 1. Spleen 2. Right Kidney 3. Left Kidney 4.Gallbladder 5.Esophagus 6. Liver 7. Stomach 8.Aorta 9. IVC 10. Portal and Splenic Veins 11. Pancreas 12 Right adrenal gland 13 Left adrenal gland.\n",
    "\n",
    "Modality: CT\n",
    "Size: 30 3D volumes (24 Training + 6 Testing)  \n",
    "Challenge: BTCV MICCAI Challenge\n",
    "\n",
    "The following figure shows image patches with the organ sub-regions that are annotated in the CT (top left) and the final labels for the whole dataset (right).\n",
    "\n",
    "Data, figures and resources are taken from: \n",
    "\n",
    "\n",
    "1. [UNETR: Transformers for 3D Medical Image Segmentation](https://arxiv.org/abs/2103.10504)\n",
    "\n",
    "2. [High-resolution 3D abdominal segmentation with random patch network fusion (MIA)](https://www.sciencedirect.com/science/article/abs/pii/S1361841520302589)\n",
    "\n",
    "3. [Efficient multi-atlas abdominal segmentation on clinically acquired CT with SIMPLE context learning (MIA)](https://www.sciencedirect.com/science/article/abs/pii/S1361841515000766?via%3Dihub)\n",
    "\n",
    "\n",
    "![image](../figures/BTCV_organs.png)\n",
    "\n",
    "\n",
    "\n",
    "The image patches show anatomies of a subject, including: \n",
    "1. large organs: spleen, liver, stomach. \n",
    "2. Smaller organs: gallbladder, esophagus, kidneys, pancreas. \n",
    "3. Vascular tissues: aorta, IVC, P&S Veins. \n",
    "4. Glands: left and right adrenal gland\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Project-MONAI/tutorials/blob/main/3d_segmentation/unetr_btcv_segmentation_3d_lightning.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python -c \"import monai\" || pip install -q \"monai-weekly[nibabel, einops]\"\n",
    "# !pip install -q pytorch-lightning~=2.0\n",
    "# !python -c \"import matplotlib\" || pip install -q matplotlib\n",
    "# %matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MONAI version: 1.4.dev2426\n",
      "Numpy version: 1.26.4\n",
      "Pytorch version: 2.3.1+cu121\n",
      "MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False\n",
      "MONAI rev id: d622a16f927841fdd7d057b7553805405f0805e4\n",
      "MONAI __file__: /home/<username>/anaconda3/envs/monai/lib/python3.9/site-packages/monai/__init__.py\n",
      "\n",
      "Optional dependencies:\n",
      "Pytorch Ignite version: 0.4.11\n",
      "ITK version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "Nibabel version: 5.2.1\n",
      "scikit-image version: 0.24.0\n",
      "scipy version: 1.13.1\n",
      "Pillow version: 10.4.0\n",
      "Tensorboard version: 2.17.0\n",
      "gdown version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "TorchVision version: 0.18.1+cu121\n",
      "tqdm version: 4.66.4\n",
      "lmdb version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "psutil version: 6.0.0\n",
      "pandas version: 2.2.2\n",
      "einops version: 0.8.0\n",
      "transformers version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "mlflow version: 2.14.2\n",
      "pynrrd version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "clearml version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "\n",
      "For details about installing the optional dependencies, please visit:\n",
      "    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from monai.losses import DiceCELoss\n",
    "from monai.inferers import sliding_window_inference\n",
    "from monai.transforms import (\n",
    "    AsDiscrete,\n",
    "    EnsureChannelFirstd,\n",
    "    Compose,\n",
    "    CropForegroundd,\n",
    "    LoadImaged,\n",
    "    Orientationd,\n",
    "    RandFlipd,\n",
    "    RandCropByPosNegLabeld,\n",
    "    RandShiftIntensityd,\n",
    "    ScaleIntensityRanged,\n",
    "    Spacingd,\n",
    "    RandRotate90d,\n",
    ")\n",
    "\n",
    "from monai.config import print_config\n",
    "from monai.metrics import DiceMetric\n",
    "from monai.networks.nets import UNETR\n",
    "\n",
    "from monai.data import (\n",
    "    DataLoader,\n",
    "    CacheDataset,\n",
    "    load_decathlon_datalist,\n",
    "    decollate_batch,\n",
    "    list_data_collate,\n",
    ")\n",
    "\n",
    "import torch\n",
    "import pytorch_lightning\n",
    "from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint\n",
    "\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.backends.cudnn.benchmark = True\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup data directory\n",
    "\n",
    "You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.  \n",
    "This allows you to save results and reuse downloads.  \n",
    "If not specified a temporary directory will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/chyang/data/dataset/project/registration/hipbone_raw_data/train/\n"
     ]
    }
   ],
   "source": [
    "os.environ['MONAI_DATA_DIRECTORY'] = '/media/chyang/data/dataset/project/registration/hipbone_raw_data/train/'\n",
    "directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n",
    "if directory is not None:\n",
    "    os.makedirs(directory, exist_ok=True)\n",
    "root_dir = tempfile.mkdtemp() if directory is None else directory\n",
    "print(root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the LightningModule (transform, network)\n",
    "The LightningModule contains a refactoring of your training code. The following module is a refactoring of the code in spleen_segmentation_3d.ipynb:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(pytorch_lightning.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self._model = UNETR(\n",
    "            in_channels=1,\n",
    "            out_channels=7,\n",
    "            img_size=(96, 96, 96),\n",
    "            feature_size=16,\n",
    "            hidden_size=768,\n",
    "            mlp_dim=3072,\n",
    "            num_heads=12,\n",
    "            pos_embed=\"perceptron\",\n",
    "            norm_name=\"instance\",\n",
    "            res_block=True,\n",
    "            conv_block=True,\n",
    "            dropout_rate=0.0,\n",
    "        ).to(device)\n",
    "\n",
    "        self.loss_function = DiceCELoss(to_onehot_y=True, softmax=True)\n",
    "        self.post_pred = AsDiscrete(argmax=True, to_onehot=14)\n",
    "        self.post_label = AsDiscrete(to_onehot=14)\n",
    "        self.dice_metric = DiceMetric(include_background=False, reduction=\"mean\", get_not_nans=False)\n",
    "        self.best_val_dice = 0\n",
    "        self.best_val_epoch = 0\n",
    "        self.max_epochs = 1300\n",
    "        self.check_val = 30\n",
    "        self.warmup_epochs = 20\n",
    "        self.metric_values = []\n",
    "        self.epoch_loss_values = []\n",
    "        self.validation_step_outputs = []\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self._model(x)\n",
    "\n",
    "    def prepare_data(self):\n",
    "        # prepare data\n",
    "        data_dir = \"/media/chyang/data/dataset/project/hipbone/nnUNet_raw/Dataset001_PelvicNAS/\"\n",
    "        split_json = \"dataset.json\"\n",
    "        datasets = data_dir + split_json\n",
    "        datalist = load_decathlon_datalist(datasets, True, \"training\")\n",
    "        val_files = load_decathlon_datalist(datasets, True, \"validation\")\n",
    "\n",
    "        train_transforms = Compose(\n",
    "            [\n",
    "                LoadImaged(keys=[\"image\", \"label\"]),\n",
    "                EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
    "                ScaleIntensityRanged(\n",
    "                    keys=[\"image\"],\n",
    "                    a_min=-175,\n",
    "                    a_max=250,\n",
    "                    b_min=0.0,\n",
    "                    b_max=1.0,\n",
    "                    clip=True,\n",
    "                ),\n",
    "                CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "                Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "                Spacingd(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    pixdim=(1.5, 1.5, 2.0),\n",
    "                    mode=(\"bilinear\", \"nearest\"),\n",
    "                ),\n",
    "                RandCropByPosNegLabeld(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    label_key=\"label\",\n",
    "                    spatial_size=(96, 96, 96),\n",
    "                    pos=1,\n",
    "                    neg=1,\n",
    "                    num_samples=4,\n",
    "                    image_key=\"image\",\n",
    "                    image_threshold=0,\n",
    "                ),\n",
    "                RandFlipd(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    spatial_axis=[0],\n",
    "                    prob=0.10,\n",
    "                ),\n",
    "                RandFlipd(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    spatial_axis=[1],\n",
    "                    prob=0.10,\n",
    "                ),\n",
    "                RandFlipd(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    spatial_axis=[2],\n",
    "                    prob=0.10,\n",
    "                ),\n",
    "                RandRotate90d(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    prob=0.10,\n",
    "                    max_k=3,\n",
    "                ),\n",
    "                RandShiftIntensityd(\n",
    "                    keys=[\"image\"],\n",
    "                    offsets=0.10,\n",
    "                    prob=0.50,\n",
    "                ),\n",
    "            ]\n",
    "        )\n",
    "        val_transforms = Compose(\n",
    "            [\n",
    "                LoadImaged(keys=[\"image\", \"label\"]),\n",
    "                EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
    "                ScaleIntensityRanged(\n",
    "                    keys=[\"image\"],\n",
    "                    a_min=-175,\n",
    "                    a_max=250,\n",
    "                    b_min=0.0,\n",
    "                    b_max=1.0,\n",
    "                    clip=True,\n",
    "                ),\n",
    "                CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "                Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "                Spacingd(\n",
    "                    keys=[\"image\", \"label\"],\n",
    "                    pixdim=(1.5, 1.5, 2.0),\n",
    "                    mode=(\"bilinear\", \"nearest\"),\n",
    "                ),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.train_ds = CacheDataset(\n",
    "            data=datalist,\n",
    "            transform=train_transforms,\n",
    "            cache_num=24,\n",
    "            cache_rate=1.0,\n",
    "            num_workers=8,\n",
    "        )\n",
    "        self.val_ds = CacheDataset(\n",
    "            data=val_files,\n",
    "            transform=val_transforms,\n",
    "            cache_num=6,\n",
    "            cache_rate=1.0,\n",
    "            num_workers=8,\n",
    "        )\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        train_loader = DataLoader(\n",
    "            self.train_ds,\n",
    "            batch_size=2,\n",
    "            shuffle=True,\n",
    "            num_workers=1,\n",
    "            pin_memory=True,\n",
    "            collate_fn=list_data_collate,\n",
    "        )\n",
    "        return train_loader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        val_loader = DataLoader(self.val_ds, batch_size=1, shuffle=False, num_workers=1, pin_memory=True)\n",
    "        return val_loader\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(self._model.parameters(), lr=1e-4, weight_decay=1e-5)\n",
    "        return optimizer\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        images, labels = (batch[\"image\"].cuda(), batch[\"label\"].cuda())\n",
    "        output = self.forward(images)\n",
    "        loss = self.loss_function(output, labels)\n",
    "        tensorboard_logs = {\"train_loss\": loss.item()}\n",
    "        return {\"loss\": loss, \"log\": tensorboard_logs}\n",
    "\n",
    "    def training_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
    "        self.epoch_loss_values.append(avg_loss.detach().cpu().numpy())\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        images, labels = batch[\"image\"], batch[\"label\"]\n",
    "        roi_size = (96, 96, 96)\n",
    "        sw_batch_size = 4\n",
    "        outputs = sliding_window_inference(images, roi_size, sw_batch_size, self.forward)\n",
    "        loss = self.loss_function(outputs, labels)\n",
    "        outputs = [self.post_pred(i) for i in decollate_batch(outputs)]\n",
    "        labels = [self.post_label(i) for i in decollate_batch(labels)]\n",
    "        self.dice_metric(y_pred=outputs, y=labels)\n",
    "        d = {\"val_loss\": loss, \"val_number\": len(outputs)}\n",
    "        self.validation_step_outputs.append(d)\n",
    "        return d\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        val_loss, num_items = 0, 0\n",
    "        for output in self.validation_step_outputs:\n",
    "            val_loss += output[\"val_loss\"].sum().item()\n",
    "            num_items += output[\"val_number\"]\n",
    "        mean_val_dice = self.dice_metric.aggregate().item()\n",
    "        self.dice_metric.reset()\n",
    "        mean_val_loss = torch.tensor(val_loss / num_items)\n",
    "        tensorboard_logs = {\n",
    "            \"val_dice\": mean_val_dice,\n",
    "            \"val_loss\": mean_val_loss,\n",
    "        }\n",
    "        if mean_val_dice > self.best_val_dice:\n",
    "            self.best_val_dice = mean_val_dice\n",
    "            self.best_val_epoch = self.current_epoch\n",
    "        print(\n",
    "            f\"current epoch: {self.current_epoch} \"\n",
    "            f\"current mean dice: {mean_val_dice:.4f}\"\n",
    "            f\"\\nbest mean dice: {self.best_val_dice:.4f} \"\n",
    "            f\"at epoch: {self.best_val_epoch}\"\n",
    "        )\n",
    "        self.metric_values.append(mean_val_dice)\n",
    "        self.validation_step_outputs.clear()  # free memory\n",
    "        return {\"log\": tensorboard_logs}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/chyang/anaconda3/envs/monai/lib/python3.9/site-packages/monai/utils/deprecate_utils.py:221: FutureWarning: monai.networks.nets.unetr UNETR.__init__:pos_embed: Argument `pos_embed` has been deprecated since version 1.2. It will be removed in version 1.4. please use `proj_type` instead.\n",
      "  warn_deprecated(argname, msg, warning_category)\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/home/chyang/anaconda3/envs/monai/lib/python3.9/site-packages/monai/utils/deprecate_utils.py:321: FutureWarning: monai.transforms.croppad.dictionary CropForegroundd.__init__:allow_smaller: Current default value of argument `allow_smaller=True` has been deprecated since version 1.2. It will be changed to `allow_smaller=False` in version 1.5.\n",
      "  warn_deprecated(argname, msg, warning_category)\n",
      "Loading dataset: 100%|██████████| 24/24 [00:51<00:00,  2.15s/it]\n",
      "Loading dataset: 100%|██████████| 6/6 [00:13<00:00,  2.20s/it]\n",
      "You are using a CUDA device ('NVIDIA GeForce RTX 3080 Ti') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
      "/home/chyang/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:613: UserWarning: Checkpoint directory /media/chyang/data/dataset/project/registration/hipbone_raw_data/train/ exists and is not empty.\n",
      "  rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name          | Type       | Params\n",
      "---------------------------------------------\n",
      "0 | _model        | UNETR      | 92.8 M\n",
      "1 | loss_function | DiceCELoss | 0     \n",
      "---------------------------------------------\n",
      "92.8 M    Trainable params\n",
      "0         Non-trainable params\n",
      "92.8 M    Total params\n",
      "371.136   Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sanity Checking: 0it [00:00, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/chyang/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sanity Checking DataLoader 0: 100%|██████████| 2/2 [00:03<00:00,  1.52s/it]current epoch: 0 current mean dice: 0.0252\n",
      "best mean dice: 0.0252 at epoch: 0\n",
      "                                                                           "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/chyang/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:   0%|          | 1/200 [00:28<1:35:14, 28.72s/it, loss=2.66, v_num=11]"
     ]
    },
    {
     "ename": "OutOfMemoryError",
     "evalue": "CUDA out of memory. Tried to allocate 432.00 MiB. GPU ",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 17\u001b[0m\n\u001b[1;32m      8\u001b[0m trainer \u001b[38;5;241m=\u001b[39m pytorch_lightning\u001b[38;5;241m.\u001b[39mTrainer(\n\u001b[1;32m      9\u001b[0m     accelerator\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgpu\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     10\u001b[0m     max_epochs\u001b[38;5;241m=\u001b[39mnet\u001b[38;5;241m.\u001b[39mmax_epochs,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     13\u001b[0m     default_root_dir\u001b[38;5;241m=\u001b[39mroot_dir,\n\u001b[1;32m     14\u001b[0m )\n\u001b[1;32m     16\u001b[0m \u001b[38;5;66;03m# train\u001b[39;00m\n\u001b[0;32m---> 17\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnet\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:608\u001b[0m, in \u001b[0;36mTrainer.fit\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    606\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_maybe_unwrap_optimized(model)\n\u001b[1;32m    607\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39m_lightning_module \u001b[38;5;241m=\u001b[39m model\n\u001b[0;32m--> 608\u001b[0m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_and_handle_interrupt\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    609\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_fit_impl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_dataloaders\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_dataloaders\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdatamodule\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mckpt_path\u001b[49m\n\u001b[1;32m    610\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py:38\u001b[0m, in \u001b[0;36m_call_and_handle_interrupt\u001b[0;34m(trainer, trainer_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m     36\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mlauncher\u001b[38;5;241m.\u001b[39mlaunch(trainer_fn, \u001b[38;5;241m*\u001b[39margs, trainer\u001b[38;5;241m=\u001b[39mtrainer, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m     37\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 38\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrainer_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     40\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _TunerExitException:\n\u001b[1;32m     41\u001b[0m     trainer\u001b[38;5;241m.\u001b[39m_call_teardown_hook()\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:650\u001b[0m, in \u001b[0;36mTrainer._fit_impl\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    643\u001b[0m ckpt_path \u001b[38;5;241m=\u001b[39m ckpt_path \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mresume_from_checkpoint\n\u001b[1;32m    644\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_ckpt_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_checkpoint_connector\u001b[38;5;241m.\u001b[39m_set_ckpt_path(\n\u001b[1;32m    645\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mfn,\n\u001b[1;32m    646\u001b[0m     ckpt_path,  \u001b[38;5;66;03m# type: ignore[arg-type]\u001b[39;00m\n\u001b[1;32m    647\u001b[0m     model_provided\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    648\u001b[0m     model_connected\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    649\u001b[0m )\n\u001b[0;32m--> 650\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mckpt_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mckpt_path\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    652\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mstopped\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1112\u001b[0m, in \u001b[0;36mTrainer._run\u001b[0;34m(self, model, ckpt_path)\u001b[0m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_checkpoint_connector\u001b[38;5;241m.\u001b[39mrestore_training_state()\n\u001b[1;32m   1110\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_checkpoint_connector\u001b[38;5;241m.\u001b[39mresume_end()\n\u001b[0;32m-> 1112\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_stage\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1114\u001b[0m log\u001b[38;5;241m.\u001b[39mdetail(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m: trainer tearing down\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   1115\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_teardown()\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1191\u001b[0m, in \u001b[0;36mTrainer._run_stage\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1189\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpredicting:\n\u001b[1;32m   1190\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_run_predict()\n\u001b[0;32m-> 1191\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_train\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1214\u001b[0m, in \u001b[0;36mTrainer._run_train\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1211\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfit_loop\u001b[38;5;241m.\u001b[39mtrainer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m   1213\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mautograd\u001b[38;5;241m.\u001b[39mset_detect_anomaly(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_detect_anomaly):\n\u001b[0;32m-> 1214\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py:199\u001b[0m, in \u001b[0;36mLoop.run\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    197\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    198\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_start(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_end()\n\u001b[1;32m    201\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py:267\u001b[0m, in \u001b[0;36mFitLoop.advance\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_data_fetcher\u001b[38;5;241m.\u001b[39msetup(dataloader, batch_to_device\u001b[38;5;241m=\u001b[39mbatch_to_device)\n\u001b[1;32m    266\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrun_training_epoch\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 267\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mepoch_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_data_fetcher\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py:199\u001b[0m, in \u001b[0;36mLoop.run\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    197\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    198\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_start(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_end()\n\u001b[1;32m    201\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py:213\u001b[0m, in \u001b[0;36mTrainingEpochLoop.advance\u001b[0;34m(self, data_fetcher)\u001b[0m\n\u001b[1;32m    210\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbatch_progress\u001b[38;5;241m.\u001b[39mincrement_started()\n\u001b[1;32m    212\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrun_training_batch\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 213\u001b[0m         batch_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbatch_progress\u001b[38;5;241m.\u001b[39mincrement_processed()\n\u001b[1;32m    217\u001b[0m \u001b[38;5;66;03m# update non-plateau LR schedulers\u001b[39;00m\n\u001b[1;32m    218\u001b[0m \u001b[38;5;66;03m# update epoch-interval ones only when we are at the end of training epoch\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py:199\u001b[0m, in \u001b[0;36mLoop.run\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    197\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    198\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_start(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_end()\n\u001b[1;32m    201\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py:88\u001b[0m, in \u001b[0;36mTrainingBatchLoop.advance\u001b[0;34m(self, kwargs)\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39mlightning_module\u001b[38;5;241m.\u001b[39mautomatic_optimization:\n\u001b[1;32m     85\u001b[0m     optimizers \u001b[38;5;241m=\u001b[39m _get_active_optimizers(\n\u001b[1;32m     86\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39moptimizers, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39moptimizer_frequencies, kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbatch_idx\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m     87\u001b[0m     )\n\u001b[0;32m---> 88\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptimizer_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43moptimizers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     90\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmanual_loop\u001b[38;5;241m.\u001b[39mrun(kwargs)\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py:199\u001b[0m, in \u001b[0;36mLoop.run\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    197\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    198\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_start(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_end()\n\u001b[1;32m    201\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:202\u001b[0m, in \u001b[0;36mOptimizerLoop.advance\u001b[0;34m(self, optimizers, kwargs)\u001b[0m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21madvance\u001b[39m(\u001b[38;5;28mself\u001b[39m, optimizers: List[Tuple[\u001b[38;5;28mint\u001b[39m, Optimizer]], kwargs: OrderedDict) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    200\u001b[0m     kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_build_kwargs(kwargs, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moptimizer_idx, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_hiddens)\n\u001b[0;32m--> 202\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_optimization\u001b[49m\u001b[43m(\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_optimizers\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptim_progress\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptimizer_position\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    203\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m result\u001b[38;5;241m.\u001b[39mloss \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    204\u001b[0m         \u001b[38;5;66;03m# automatic optimization assumes a loss needs to be returned for extras to be considered as the batch\u001b[39;00m\n\u001b[1;32m    205\u001b[0m         \u001b[38;5;66;03m# would be skipped otherwise\u001b[39;00m\n\u001b[1;32m    206\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_outputs[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moptimizer_idx] \u001b[38;5;241m=\u001b[39m result\u001b[38;5;241m.\u001b[39masdict()\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:249\u001b[0m, in \u001b[0;36mOptimizerLoop._run_optimization\u001b[0;34m(self, kwargs, optimizer)\u001b[0m\n\u001b[1;32m    241\u001b[0m         closure()\n\u001b[1;32m    243\u001b[0m \u001b[38;5;66;03m# ------------------------------\u001b[39;00m\n\u001b[1;32m    244\u001b[0m \u001b[38;5;66;03m# BACKWARD PASS\u001b[39;00m\n\u001b[1;32m    245\u001b[0m \u001b[38;5;66;03m# ------------------------------\u001b[39;00m\n\u001b[1;32m    246\u001b[0m \u001b[38;5;66;03m# gradient update with accumulated gradients\u001b[39;00m\n\u001b[1;32m    247\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    248\u001b[0m     \u001b[38;5;66;03m# the `batch_idx` is optional with inter-batch parallelism\u001b[39;00m\n\u001b[0;32m--> 249\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_optimizer_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mopt_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbatch_idx\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    251\u001b[0m result \u001b[38;5;241m=\u001b[39m closure\u001b[38;5;241m.\u001b[39mconsume_result()\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m result\u001b[38;5;241m.\u001b[39mloss \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    254\u001b[0m     \u001b[38;5;66;03m# if no result, user decided to skip optimization\u001b[39;00m\n\u001b[1;32m    255\u001b[0m     \u001b[38;5;66;03m# otherwise update running loss + reset accumulated loss\u001b[39;00m\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;66;03m# TODO: find proper way to handle updating running loss\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:370\u001b[0m, in \u001b[0;36mOptimizerLoop._optimizer_step\u001b[0;34m(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)\u001b[0m\n\u001b[1;32m    362\u001b[0m     rank_zero_deprecation(\n\u001b[1;32m    363\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe NVIDIA/apex AMP implementation has been deprecated upstream. Consequently, its integration inside\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    364\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m PyTorch Lightning has been deprecated in v1.9.0 and will be removed in v2.0.0.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    367\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m return True.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    368\u001b[0m     )\n\u001b[1;32m    369\u001b[0m     kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124musing_native_amp\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39mprecision_plugin, MixedPrecisionPlugin)\n\u001b[0;32m--> 370\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_lightning_module_hook\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    371\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43moptimizer_step\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    372\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcurrent_epoch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    373\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbatch_idx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    374\u001b[0m \u001b[43m    \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    375\u001b[0m \u001b[43m    \u001b[49m\u001b[43mopt_idx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    376\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtrain_step_and_backward_closure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    377\u001b[0m \u001b[43m    \u001b[49m\u001b[43mon_tpu\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43misinstance\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maccelerator\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mTPUAccelerator\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    378\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m    379\u001b[0m \u001b[43m    \u001b[49m\u001b[43musing_lbfgs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_lbfgs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    380\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m should_accumulate:\n\u001b[1;32m    383\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moptim_progress\u001b[38;5;241m.\u001b[39moptimizer\u001b[38;5;241m.\u001b[39mstep\u001b[38;5;241m.\u001b[39mincrement_completed()\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1356\u001b[0m, in \u001b[0;36mTrainer._call_lightning_module_hook\u001b[0;34m(self, hook_name, pl_module, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1353\u001b[0m pl_module\u001b[38;5;241m.\u001b[39m_current_fx_name \u001b[38;5;241m=\u001b[39m hook_name\n\u001b[1;32m   1355\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[LightningModule]\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpl_module\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mhook_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m-> 1356\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1358\u001b[0m \u001b[38;5;66;03m# restore current_fx when nested context\u001b[39;00m\n\u001b[1;32m   1359\u001b[0m pl_module\u001b[38;5;241m.\u001b[39m_current_fx_name \u001b[38;5;241m=\u001b[39m prev_fx_name\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/core/module.py:1742\u001b[0m, in \u001b[0;36mLightningModule.optimizer_step\u001b[0;34m(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_lbfgs)\u001b[0m\n\u001b[1;32m   1663\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21moptimizer_step\u001b[39m(\n\u001b[1;32m   1664\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m   1665\u001b[0m     epoch: \u001b[38;5;28mint\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1671\u001b[0m     using_lbfgs: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m   1672\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1673\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   1674\u001b[0m \u001b[38;5;124;03m    Override this method to adjust the default way the :class:`~pytorch_lightning.trainer.trainer.Trainer` calls\u001b[39;00m\n\u001b[1;32m   1675\u001b[0m \u001b[38;5;124;03m    each optimizer.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1740\u001b[0m \n\u001b[1;32m   1741\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m-> 1742\u001b[0m     \u001b[43moptimizer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclosure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moptimizer_closure\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py:169\u001b[0m, in \u001b[0;36mLightningOptimizer.step\u001b[0;34m(self, closure, **kwargs)\u001b[0m\n\u001b[1;32m    166\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m MisconfigurationException(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mWhen `optimizer.step(closure)` is called, the closure should be callable\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    168\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_strategy \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 169\u001b[0m step_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_strategy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptimizer_step\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_optimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_optimizer_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    171\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_on_after_step()\n\u001b[1;32m    173\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m step_output\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py:234\u001b[0m, in \u001b[0;36mStrategy.optimizer_step\u001b[0;34m(self, optimizer, opt_idx, closure, model, **kwargs)\u001b[0m\n\u001b[1;32m    232\u001b[0m \u001b[38;5;66;03m# TODO(fabric): remove assertion once strategy's optimizer_step typing is fixed\u001b[39;00m\n\u001b[1;32m    233\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(model, pl\u001b[38;5;241m.\u001b[39mLightningModule)\n\u001b[0;32m--> 234\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprecision_plugin\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptimizer_step\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    235\u001b[0m \u001b[43m    \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer_idx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mopt_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    236\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:119\u001b[0m, in \u001b[0;36mPrecisionPlugin.optimizer_step\u001b[0;34m(self, optimizer, model, optimizer_idx, closure, **kwargs)\u001b[0m\n\u001b[1;32m    117\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Hook to run the optimizer step.\"\"\"\u001b[39;00m\n\u001b[1;32m    118\u001b[0m closure \u001b[38;5;241m=\u001b[39m partial(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_wrap_closure, model, optimizer, optimizer_idx, closure)\n\u001b[0;32m--> 119\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43moptimizer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclosure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/optim/optimizer.py:391\u001b[0m, in \u001b[0;36mOptimizer.profile_hook_step.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    386\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    387\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[1;32m    388\u001b[0m                 \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfunc\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m must return None or a tuple of (new_args, new_kwargs), but got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mresult\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    389\u001b[0m             )\n\u001b[0;32m--> 391\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    392\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_optimizer_step_code()\n\u001b[1;32m    394\u001b[0m \u001b[38;5;66;03m# call optimizer step post hooks\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/optim/optimizer.py:76\u001b[0m, in \u001b[0;36m_use_grad_for_differentiable.<locals>._use_grad\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     74\u001b[0m     torch\u001b[38;5;241m.\u001b[39mset_grad_enabled(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdefaults[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdifferentiable\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m     75\u001b[0m     torch\u001b[38;5;241m.\u001b[39m_dynamo\u001b[38;5;241m.\u001b[39mgraph_break()\n\u001b[0;32m---> 76\u001b[0m     ret \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     78\u001b[0m     torch\u001b[38;5;241m.\u001b[39m_dynamo\u001b[38;5;241m.\u001b[39mgraph_break()\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/optim/adamw.py:165\u001b[0m, in \u001b[0;36mAdamW.step\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m    163\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m closure \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    164\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39menable_grad():\n\u001b[0;32m--> 165\u001b[0m         loss \u001b[38;5;241m=\u001b[39m \u001b[43mclosure\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    167\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m group \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparam_groups:\n\u001b[1;32m    168\u001b[0m     params_with_grad \u001b[38;5;241m=\u001b[39m []\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:105\u001b[0m, in \u001b[0;36mPrecisionPlugin._wrap_closure\u001b[0;34m(self, model, optimizer, optimizer_idx, closure)\u001b[0m\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_wrap_closure\u001b[39m(\n\u001b[1;32m     93\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     94\u001b[0m     model: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpl.LightningModule\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     97\u001b[0m     closure: Callable[[], Any],\n\u001b[1;32m     98\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     99\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"This double-closure allows makes sure the ``closure`` is executed before the\u001b[39;00m\n\u001b[1;32m    100\u001b[0m \u001b[38;5;124;03m    ``on_before_optimizer_step`` hook is called.\u001b[39;00m\n\u001b[1;32m    101\u001b[0m \n\u001b[1;32m    102\u001b[0m \u001b[38;5;124;03m    The closure (generally) runs ``backward`` so this allows inspecting gradients in this hook. This structure is\u001b[39;00m\n\u001b[1;32m    103\u001b[0m \u001b[38;5;124;03m    consistent with the ``PrecisionPlugin`` subclasses that cannot pass ``optimizer.step(closure)`` directly.\u001b[39;00m\n\u001b[1;32m    104\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 105\u001b[0m     closure_result \u001b[38;5;241m=\u001b[39m \u001b[43mclosure\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    106\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_after_closure(model, optimizer, optimizer_idx)\n\u001b[1;32m    107\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m closure_result\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:149\u001b[0m, in \u001b[0;36mClosure.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    148\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs: Any, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Optional[Tensor]:\n\u001b[0;32m--> 149\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclosure\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    150\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_result\u001b[38;5;241m.\u001b[39mloss\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:144\u001b[0m, in \u001b[0;36mClosure.closure\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    141\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_zero_grad_fn()\n\u001b[1;32m    143\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_fn \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m step_output\u001b[38;5;241m.\u001b[39mclosure_loss \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 144\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_backward_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstep_output\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclosure_loss\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    146\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m step_output\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:305\u001b[0m, in \u001b[0;36mOptimizerLoop._make_backward_fn.<locals>.backward_fn\u001b[0;34m(loss)\u001b[0m\n\u001b[1;32m    304\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbackward_fn\u001b[39m(loss: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 305\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_strategy_hook\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbackward\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mloss\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mopt_idx\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1494\u001b[0m, in \u001b[0;36mTrainer._call_strategy_hook\u001b[0;34m(self, hook_name, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1491\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m   1493\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[Strategy]\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mhook_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m-> 1494\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# restore current_fx when nested context\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m pl_module\u001b[38;5;241m.\u001b[39m_current_fx_name \u001b[38;5;241m=\u001b[39m prev_fx_name\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py:207\u001b[0m, in \u001b[0;36mStrategy.backward\u001b[0;34m(self, closure_loss, optimizer, optimizer_idx, *args, **kwargs)\u001b[0m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    205\u001b[0m closure_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprecision_plugin\u001b[38;5;241m.\u001b[39mpre_backward(closure_loss, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module)\n\u001b[0;32m--> 207\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprecision_plugin\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclosure_loss\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlightning_module\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    209\u001b[0m closure_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprecision_plugin\u001b[38;5;241m.\u001b[39mpost_backward(closure_loss, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module)\n\u001b[1;32m    210\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpost_backward(closure_loss)\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:67\u001b[0m, in \u001b[0;36mPrecisionPlugin.backward\u001b[0;34m(self, tensor, model, optimizer, optimizer_idx, *args, **kwargs)\u001b[0m\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbackward\u001b[39m(  \u001b[38;5;66;03m# type: ignore[override]\u001b[39;00m\n\u001b[1;32m     48\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     49\u001b[0m     tensor: Tensor,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     54\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m     55\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     56\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Performs the actual backpropagation.\u001b[39;00m\n\u001b[1;32m     57\u001b[0m \n\u001b[1;32m     58\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     65\u001b[0m \u001b[38;5;124;03m        \\**kwargs: Keyword arguments for the same purpose as ``*args``.\u001b[39;00m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 67\u001b[0m     \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtensor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/pytorch_lightning/core/module.py:1486\u001b[0m, in \u001b[0;36mLightningModule.backward\u001b[0;34m(self, loss, optimizer, optimizer_idx, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1484\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_fabric\u001b[38;5;241m.\u001b[39mbackward(loss, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m   1485\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1486\u001b[0m     \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/_tensor.py:516\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    469\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Computes the gradient of current tensor wrt graph leaves.\u001b[39;00m\n\u001b[1;32m    470\u001b[0m \n\u001b[1;32m    471\u001b[0m \u001b[38;5;124;03mThe graph is differentiated using the chain rule. If the tensor is\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    513\u001b[0m \u001b[38;5;124;03m        used to compute the attr::tensors.\u001b[39;00m\n\u001b[1;32m    514\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    515\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m--> 516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhandle_torch_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    517\u001b[0m \u001b[43m        \u001b[49m\u001b[43mTensor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    518\u001b[0m \u001b[43m        \u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    519\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    520\u001b[0m \u001b[43m        \u001b[49m\u001b[43mgradient\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    521\u001b[0m \u001b[43m        \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    522\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    523\u001b[0m \u001b[43m        \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    524\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    525\u001b[0m torch\u001b[38;5;241m.\u001b[39mautograd\u001b[38;5;241m.\u001b[39mbackward(\n\u001b[1;32m    526\u001b[0m     \u001b[38;5;28mself\u001b[39m, gradient, retain_graph, create_graph, inputs\u001b[38;5;241m=\u001b[39minputs\n\u001b[1;32m    527\u001b[0m )\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/overrides.py:1636\u001b[0m, in \u001b[0;36mhandle_torch_function\u001b[0;34m(public_api, relevant_args, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1630\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDefining your `__torch_function__ as a plain method is deprecated and \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1631\u001b[0m                   \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwill be an error in future, please define it as a classmethod.\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   1632\u001b[0m                   \u001b[38;5;167;01mDeprecationWarning\u001b[39;00m)\n\u001b[1;32m   1634\u001b[0m \u001b[38;5;66;03m# Use `public_api` instead of `implementation` so __torch_function__\u001b[39;00m\n\u001b[1;32m   1635\u001b[0m \u001b[38;5;66;03m# implementations can do equality/identity comparisons.\u001b[39;00m\n\u001b[0;32m-> 1636\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mtorch_func_method\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpublic_api\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtypes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1638\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m result \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mNotImplemented\u001b[39m:\n\u001b[1;32m   1639\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/data/meta_tensor.py:282\u001b[0m, in \u001b[0;36mMetaTensor.__torch_function__\u001b[0;34m(cls, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m    280\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m kwargs \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    281\u001b[0m     kwargs \u001b[38;5;241m=\u001b[39m {}\n\u001b[0;32m--> 282\u001b[0m ret \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__torch_function__\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtypes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    283\u001b[0m \u001b[38;5;66;03m# if `out` has been used as argument, metadata is not copied, nothing to do.\u001b[39;00m\n\u001b[1;32m    284\u001b[0m \u001b[38;5;66;03m# if \"out\" in kwargs:\u001b[39;00m\n\u001b[1;32m    285\u001b[0m \u001b[38;5;66;03m#     return ret\u001b[39;00m\n\u001b[1;32m    286\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _not_requiring_metadata(ret):\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/_tensor.py:1443\u001b[0m, in \u001b[0;36mTensor.__torch_function__\u001b[0;34m(cls, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   1440\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mNotImplemented\u001b[39m\n\u001b[1;32m   1442\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m _C\u001b[38;5;241m.\u001b[39mDisableTorchFunctionSubclass():\n\u001b[0;32m-> 1443\u001b[0m     ret \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1444\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m func \u001b[38;5;129;01min\u001b[39;00m get_default_nowrap_functions():\n\u001b[1;32m   1445\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m ret\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/_tensor.py:525\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    515\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    517\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    518\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    523\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    524\u001b[0m     )\n\u001b[0;32m--> 525\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    526\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    527\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/autograd/__init__.py:267\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    262\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    264\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    266\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 267\u001b[0m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    268\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    269\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    270\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    271\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    272\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    273\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    274\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/autograd/graph.py:744\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    742\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[1;32m    743\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 744\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    745\u001b[0m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    746\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[1;32m    747\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    748\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 432.00 MiB. GPU "
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m在当前单元格或上一个单元格中执行代码时 Kernel 崩溃。\n",
      "\u001b[1;31m请查看单元格中的代码，以确定故障的可能原因。\n",
      "\u001b[1;31m单击<a href='https://aka.ms/vscodeJupyterKernelCrash'>此处</a>了解详细信息。\n",
      "\u001b[1;31m有关更多详细信息，请查看 Jupyter <a href='command:jupyter.viewOutput'>log</a>。"
     ]
    }
   ],
   "source": [
    "# initialise the LightningModule\n",
    "net = Net()\n",
    "\n",
    "# set up checkpoints\n",
    "checkpoint_callback = ModelCheckpoint(dirpath=root_dir, filename=\"best_metric_model\")\n",
    "\n",
    "# initialise Lightning's trainer.\n",
    "trainer = pytorch_lightning.Trainer(\n",
    "    accelerator=\"gpu\",\n",
    "    max_epochs=net.max_epochs,\n",
    "    check_val_every_n_epoch=net.check_val,\n",
    "    callbacks=checkpoint_callback,\n",
    "    default_root_dir=root_dir,\n",
    ")\n",
    "\n",
    "# train\n",
    "trainer.fit(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the loss and metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "eval_num = 250\n",
    "plt.figure(\"train\", (12, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"Iteration Average Loss\")\n",
    "x = [eval_num * (i + 1) for i in range(len(net.epoch_loss_values))]\n",
    "y = net.epoch_loss_values\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.plot(x, y)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Val Mean Dice\")\n",
    "x = [eval_num * (i + 1) for i in range(len(net.metric_values))]\n",
    "y = net.metric_values\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.plot(x, y)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check best model output with the input image and label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_map = {\n",
    "    \"img0035.nii.gz\": 170,\n",
    "    \"img0036.nii.gz\": 230,\n",
    "    \"img0037.nii.gz\": 204,\n",
    "    \"img0038.nii.gz\": 204,\n",
    "    \"img0039.nii.gz\": 204,\n",
    "    \"img0040.nii.gz\": 180,\n",
    "}\n",
    "case_num = 4\n",
    "net.load_from_checkpoint(os.path.join(root_dir, \"best_metric_model-v1.ckpt\"))\n",
    "net.eval()\n",
    "net.to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    img_name = os.path.split(net.val_ds[case_num][\"image\"].meta[\"filename_or_obj\"])[1]\n",
    "    img = net.val_ds[case_num][\"image\"]\n",
    "    label = net.val_ds[case_num][\"label\"]\n",
    "    val_inputs = torch.unsqueeze(img, 1).cuda()\n",
    "    val_labels = torch.unsqueeze(label, 1).cuda()\n",
    "    val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, net, overlap=0.8)\n",
    "    plt.figure(\"check\", (18, 6))\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.title(\"image\")\n",
    "    plt.imshow(val_inputs.cpu().numpy()[0, 0, :, :, slice_map[img_name]], cmap=\"gray\")\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.title(\"label\")\n",
    "    plt.imshow(val_labels.cpu().numpy()[0, 0, :, :, slice_map[img_name]])\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.title(\"output\")\n",
    "    plt.imshow(torch.argmax(val_outputs, dim=1).detach().cpu()[0, :, :, slice_map[img_name]])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup data directory\n",
    "\n",
    "Remove directory if a temporary was used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if directory is None:\n",
    "    shutil.rmtree(root_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
